// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/expression/vector_function.h>

namespace kumo::pollux::functions {
    namespace detail {
        template<typename T, bool kIsConstant>
        inline auto loadSimdData(const T *rawData, vector_size_t offset) {
            using d_type = ksimd::batch<T>;
            if constexpr (kIsConstant) {
                return ksimd::broadcast<T>(rawData[0]);
            }
            return d_type::load_unaligned(rawData + offset);
        }

        inline uint64_t to64Bits(const int8_t *resultData) {
            using d_type = ksimd::batch<int8_t>;
            constexpr auto numScalarElements = d_type::size;
            static_assert(
                numScalarElements == 16 || numScalarElements == 32 ||
                numScalarElements == 64,
                "Unsupported number of scalar elements");
            uint64_t res = 0UL;
            if constexpr (numScalarElements == 64) {
                res = simd::toBitMask(ksimd::batch_bool<int8_t>(
                    simd::reinterpretBatch<uint8_t>(d_type::load_unaligned(resultData))));
            } else if constexpr (numScalarElements == 32) {
                auto *addr = reinterpret_cast<uint32_t *>(&res);
                *(addr) = simd::toBitMask(ksimd::batch_bool<int8_t>(
                    simd::reinterpretBatch<uint8_t>(d_type::load_unaligned(resultData))));
                *(addr + 1) = simd::toBitMask(
                    ksimd::batch_bool<int8_t>(simd::reinterpretBatch<uint8_t>(
                        d_type::load_unaligned(resultData + 32))));
            } else if constexpr (numScalarElements == 16) {
                auto *addr = reinterpret_cast<uint16_t *>(&res);
                *(addr) = simd::toBitMask(ksimd::batch_bool<int8_t>(
                    simd::reinterpretBatch<uint8_t>(d_type::load_unaligned(resultData))));
                *(addr + 1) = simd::toBitMask(
                    ksimd::batch_bool<int8_t>(simd::reinterpretBatch<uint8_t>(
                        d_type::load_unaligned(resultData + 16))));
                *(addr + 2) = simd::toBitMask(
                    ksimd::batch_bool<int8_t>(simd::reinterpretBatch<uint8_t>(
                        d_type::load_unaligned(resultData + 32))));
                *(addr + 3) = simd::toBitMask(
                    ksimd::batch_bool<int8_t>(simd::reinterpretBatch<uint8_t>(
                        d_type::load_unaligned(resultData + 48))));
            }
            return res;
        }

        template<typename A, typename B, typename Compare>
        void applyAutoSimdComparisonInternal(
            const SelectivityVector &rows,
            const A *__restrict rawA,
            const B *__restrict rawB,
            Compare cmp,
            VectorPtr &result) {
            int8_t tempBuffer[64];
            int8_t *__restrict resultData = tempBuffer;
            const vector_size_t rowsBegin = rows.begin();
            const vector_size_t rowsEnd = rows.end();
            auto *rowsData = reinterpret_cast<const uint64_t *>(rows.allBits());
            auto *resultVector = result->as_unchecked<FlatVector<bool> >();
            auto *rawResult = resultVector->mutableRawValues<uint64_t>();
            if (rows.isAllSelected()) {
                auto i = 0;
                for (; i + 64 <= rowsEnd; i += 64) {
                    for (auto j = 0; j < 64; ++j) {
                        resultData[j] = cmp(rawA, rawB, i + j) ? -1 : 0;
                    }
                    rawResult[i / 64] = to64Bits(resultData);
                }
                for (; i < rowsEnd; ++i) {
                    bits::setBit(rawResult, i, cmp(rawA, rawB, i));
                }
            } else {
                static constexpr uint64_t kAllSet = -1ULL;
                bits::forEachWord(
                    rowsBegin,
                    rowsEnd,
                    [&](int32_t idx, uint64_t mask) {
                        auto word = rowsData[idx] & mask;
                        if (!word) {
                            return;
                        }
                        const size_t start = idx * 64;
                        while (word) {
                            auto index = start + __builtin_ctzll(word);
                            bits::setBit(rawResult, index, cmp(rawA, rawB, index));
                            word &= word - 1;
                        }
                    },
                    [&](int32_t idx) {
                        auto word = rowsData[idx];
                        const size_t start = idx * 64;
                        if (kAllSet == word) {
                            // Do 64 comparisons in a batch, set results by SIMD.
                            for (size_t row = 0; row < 64; ++row) {
                                resultData[row] = cmp(rawA, rawB, row + start) ? -1 : 0;
                            }
                            rawResult[idx] = to64Bits(resultData);
                        } else {
                            while (word) {
                                auto index = __builtin_ctzll(word);
                                resultData[index] = cmp(rawA, rawB, start + index) ? -1 : 0;
                                word &= word - 1;
                            }
                            // Set results only for selected rows.
                            uint64_t mask = rowsData[idx];
                            rawResult[idx] =
                                    (rawResult[idx] & ~mask) | (to64Bits(resultData) & mask);
                        }
                    });
            }
        }

        inline bool isDictEncoding(const VectorPtr &arg) {
            return arg->encoding() == VectorEncoding::Simple::DICTIONARY;
        }
    } // namespace detail

    template<
        typename T,
        bool kIsLeftConstant,
        bool kIsRightConstant,
        typename ComparisonOp>
    void applySimdComparison(
        const vector_size_t begin,
        const vector_size_t end,
        const T *rawLhs,
        const T *rawRhs,
        uint8_t *rawResult) {
        using d_type = ksimd::batch<T>;
        constexpr auto numScalarElements = d_type::size;
        const auto vectorEnd = (end - begin) - (end - begin) % numScalarElements;
        static_assert(
            numScalarElements == 2 || numScalarElements == 4 ||
            numScalarElements == 8 || numScalarElements == 16 ||
            numScalarElements == 32,
            "Unsupported number of scalar elements");
        if constexpr (numScalarElements == 2 || numScalarElements == 4) {
            for (auto i = begin; i < vectorEnd; i += 8) {
                rawResult[i / 8] = 0;
                for (auto j = 0; j < 8 && (i + j) < vectorEnd; j += numScalarElements) {
                    auto left = detail::loadSimdData<T, kIsLeftConstant>(rawLhs, i + j);
                    auto right = detail::loadSimdData<T, kIsRightConstant>(rawRhs, i + j);

                    uint8_t res = simd::toBitMask(ComparisonOp()(left, right));
                    rawResult[i / 8] |= res << j;
                }
            }
        } else {
            for (auto i = begin; i < vectorEnd; i += numScalarElements) {
                auto left = detail::loadSimdData<T, kIsLeftConstant>(rawLhs, i);
                auto right = detail::loadSimdData<T, kIsRightConstant>(rawRhs, i);

                auto res = simd::toBitMask(ComparisonOp()(left, right));
                if constexpr (numScalarElements == 8) {
                    rawResult[i / 8] = res;
                } else if constexpr (numScalarElements == 16) {
                    uint16_t *addr = reinterpret_cast<uint16_t *>(rawResult + i / 8);
                    *addr = res;
                } else if constexpr (numScalarElements == 32) {
                    uint32_t *addr = reinterpret_cast<uint32_t *>(rawResult + i / 8);
                    *addr = res;
                }
            }
        }

        // Evaluate remaining values.
        for (auto i = vectorEnd; i < end; i++) {
            if constexpr (kIsRightConstant && kIsLeftConstant) {
                bits::setBit(rawResult, i, ComparisonOp()(rawLhs[0], rawRhs[0]));
            } else if constexpr (kIsRightConstant) {
                bits::setBit(rawResult, i, ComparisonOp()(rawLhs[i], rawRhs[0]));
            } else if constexpr (kIsLeftConstant) {
                bits::setBit(rawResult, i, ComparisonOp()(rawLhs[0], rawRhs[i]));
            } else {
                bits::setBit(rawResult, i, ComparisonOp()(rawLhs[i], rawRhs[i]));
            }
        }
    }

    template<typename T, typename ComparisonOp>
    void applySimdComparison(
        const SelectivityVector &rows,
        std::vector<VectorPtr> &args,
        VectorPtr &result) {
        auto resultVector = result->as_unchecked<FlatVector<bool> >();
        auto rawResult = resultVector->mutableRawValues<uint8_t>();
        if (args[0]->is_constant_encoding() && args[1]->is_constant_encoding()) {
            auto l = args[0]->as_unchecked<ConstantVector<T> >()->value_at(0);
            auto r = args[1]->as_unchecked<ConstantVector<T> >()->value_at(0);
            applySimdComparison<T, true, true, ComparisonOp>(
                rows.begin(), rows.end(), &l, &r, rawResult);
        } else if (args[0]->is_constant_encoding()) {
            auto l = args[0]->as_unchecked<ConstantVector<T> >()->value_at(0);
            auto rawRhs = args[1]->as_unchecked<FlatVector<T> >()->rawValues();
            applySimdComparison<T, true, false, ComparisonOp>(
                rows.begin(), rows.end(), &l, rawRhs, rawResult);
        } else if (args[1]->is_constant_encoding()) {
            auto rawLhs = args[0]->as_unchecked<FlatVector<T> >()->rawValues();
            auto r = args[1]->as_unchecked<ConstantVector<T> >()->value_at(0);
            applySimdComparison<T, false, true, ComparisonOp>(
                rows.begin(), rows.end(), rawLhs, &r, rawResult);
        } else {
            auto rawLhs = args[0]->as_unchecked<FlatVector<T> >()->rawValues();
            auto rawRhs = args[1]->as_unchecked<FlatVector<T> >()->rawValues();
            applySimdComparison<T, false, false, ComparisonOp>(
                rows.begin(), rows.end(), rawLhs, rawRhs, rawResult);
        }
    }

    template<typename A, typename B, typename Compare, typename... Args>
    void applyAutoSimdComparison(
        const SelectivityVector &rows,
        std::vector<VectorPtr> &args,
        exec::EvalCtx &context,
        VectorPtr &result,
        Args... cmpArgs) {
        const Compare cmp;
        if (args[0]->is_flat_encoding() && args[1]->is_flat_encoding()) {
            const A *__restrict rawA =
                    args[0]->as_unchecked<FlatVector<A> >()->template rawValues<A>();
            const B *__restrict rawB =
                    args[1]->as_unchecked<FlatVector<B> >()->template rawValues<B>();
            detail::applyAutoSimdComparisonInternal(
                rows,
                rawA,
                rawB,
                [&](const A *__restrict rawA, const B *__restrict rawB, int i) {
                    if constexpr (sizeof...(cmpArgs) > 0) {
                        return Compare::apply(rawA[i], rawB[i], cmpArgs...);
                    } else {
                        return cmp(rawA[i], rawB[i]);
                    }
                },
                result);
        } else if (args[0]->is_constant_encoding() && args[1]->is_flat_encoding()) {
            const A constA = args[0]->as_unchecked<ConstantVector<A> >()->value_at(0);
            const A *__restrict rawA = &constA;
            const B *__restrict rawB =
                    args[1]->as_unchecked<FlatVector<B> >()->template rawValues<B>();
            detail::applyAutoSimdComparisonInternal(
                rows,
                rawA,
                rawB,
                [&](const A *__restrict rawA, const B *__restrict rawB, int i) {
                    if constexpr (sizeof...(cmpArgs) > 0) {
                        return Compare::apply(rawA[0], rawB[i], cmpArgs...);
                    } else {
                        return cmp(rawA[0], rawB[i]);
                    }
                },
                result);
        } else if (args[0]->is_flat_encoding() && args[1]->is_constant_encoding()) {
            const A *__restrict rawA =
                    args[0]->as_unchecked<FlatVector<A> >()->template rawValues<A>();
            const B constB = args[1]->as_unchecked<ConstantVector<B> >()->value_at(0);
            const B *__restrict rawB = &constB;
            detail::applyAutoSimdComparisonInternal(
                rows,
                rawA,
                rawB,
                [&](const A *__restrict rawA, const B *__restrict rawB, int i) {
                    if constexpr (sizeof...(cmpArgs) > 0) {
                        return Compare::apply(rawA[i], rawB[0], cmpArgs...);
                    } else {
                        return cmp(rawA[i], rawB[0]);
                    }
                },
                result);
        } else if (args[0]->is_constant_encoding() && args[1]->is_constant_encoding()) {
            const A constA = args[0]->as_unchecked<ConstantVector<A> >()->value_at(0);
            const A *__restrict rawA = &constA;
            const B constB = args[1]->as_unchecked<ConstantVector<B> >()->value_at(0);
            const B *__restrict rawB = &constB;
            detail::applyAutoSimdComparisonInternal(
                rows,
                rawA,
                rawB,
                [&](const A *__restrict rawA, const B *__restrict rawB, int i) {
                    if constexpr (sizeof...(cmpArgs) > 0) {
                        return Compare::apply(rawA[0], rawB[0], cmpArgs...);
                    } else {
                        return cmp(rawA[0], rawB[0]);
                    }
                },
                result);
        } else if (args[0]->is_flat_encoding() && detail::isDictEncoding(args[1])) {
            const A *__restrict rawA =
                    args[0]->as_unchecked<FlatVector<A> >()->rawValues();
            exec::LocalDecodedVector localBDecoded(context, *args[1], rows);
            const B *__restrict rawB = localBDecoded.get()->data<B>();
            const vector_size_t *__restrict indexB = localBDecoded.get()->indices();
            detail::applyAutoSimdComparisonInternal(
                rows,
                rawA,
                rawB,
                [&](const A *__restrict rawA, const B *__restrict rawB, int i) {
                    if constexpr (sizeof...(cmpArgs) > 0) {
                        return Compare::apply(rawA[i], rawB[indexB[i]], cmpArgs...);
                    } else {
                        return cmp(rawA[i], rawB[indexB[i]]);
                    }
                },
                result);
        } else if (detail::isDictEncoding(args[0]) && args[1]->is_flat_encoding()) {
            exec::LocalDecodedVector localADecoded(context, *args[0], rows);
            const A *__restrict rawA = localADecoded.get()->data<A>();
            const vector_size_t *__restrict indexA = localADecoded.get()->indices();
            const B *__restrict rawB =
                    args[1]->as_unchecked<FlatVector<B> >()->rawValues();
            detail::applyAutoSimdComparisonInternal(
                rows,
                rawA,
                rawB,
                [&](const A *__restrict rawA, const B *__restrict rawB, int i) {
                    if constexpr (sizeof...(cmpArgs) > 0) {
                        return Compare::apply(rawA[indexA[i]], rawB[i], cmpArgs...);
                    } else {
                        return cmp(rawA[indexA[i]], rawB[i]);
                    }
                },
                result);
        } else {
            POLLUX_UNREACHABLE();
        }
    }

    template<typename A, typename B>
    bool shouldApplyAutoSimdComparison(
        const SelectivityVector &rows,
        std::vector<VectorPtr> &args) {
        if (rows.end() - rows.begin() > 64) {
            if ((args[0]->is_flat_encoding() || args[0]->is_constant_encoding()) &&
                (args[1]->is_flat_encoding() || args[1]->is_constant_encoding())) {
                return true;
            } else if (args[0]->is_flat_encoding() && detail::isDictEncoding(args[1])) {
                return true;
            } else if (detail::isDictEncoding(args[0]) && args[1]->is_flat_encoding()) {
                return true;
            }
        }
        return false;
    }
} // namespace kumo::pollux::functions
